# -*- coding: utf-8 -*-


from __future__ import print_function

import sys
import io
import collections
import pickle
import os.path

# -------------------------------数据预处理---------------------------#

datdir = '/Users/vista/PycharmProjects/data/poetry/'
modeldir = '/Users/vista/PycharmProjects/data/model/poetry/'

if( os.path.exists(datdir) == False):
    datdir = '/opt/vista/data/poetry/'
if( os.path.exists(modeldir) == False ):
    modeldir = '/opt/vista/data/model/poetry/'

poetry_file = datdir+'poetry.txt'

reload(sys)
sys.setdefaultencoding('utf8')


def parse_data():
    # 诗集
    poetrys = []
    with io.open(poetry_file, "r", encoding='utf-8', ) as f:
        for line in f:
            try:
                title, content = line.strip().split(':')
                content = content.replace(' ', '')
                if '_' in content or '(' in content or '（' in content or '《' in content or '[' in content:
                    continue
                if len(content) < 5 or len(content) > 79:
                    continue
                content = '[' + content + ']'
                poetrys.append(content)
            except Exception as e:
                pass

    # 按诗的字数排序
    poetrys = sorted(poetrys, key=lambda line: len(line))
    print('唐诗总数: ', len(poetrys))

    # 统计每个字出现次数
    all_words = []
    for poetry in poetrys:
        all_words += [word for word in poetry]
    counter = collections.Counter(all_words)
    count_pairs = sorted(counter.items(), key=lambda x: -x[1])
    words, _ = zip(*count_pairs)

    words = words[:len(words)] + (' ',)
    word_num_map = dict(zip(words, range(len(words))))
    to_num = lambda word: word_num_map.get(word, len(words))
    poetrys_vector = [list(map(to_num, poetry)) for poetry in poetrys]

    return poetrys_vector,words,word_num_map

def dump_data(pv,len):
    f = io.open(datdir+'poetry.pickle','wb')
    pickle.dump({'pv': pv,'len': len},f)
    f.close()

def load_data():
    f = io.open(datdir+'poetry.pickle','rb')
    dat = pickle.load(f)
    f.close()
    return dat

# pv,len = parse_data()
# dump_data(pv,len)


